
import numpy as np
import os


def vector(filename):
    vector = np.zeros(32 * 32)
    index = 0
    with open(filename, "r", encoding="utf-8") as fr:
        for i in range(32):
            line = fr.readline()
            for j in range(32):
                vector[index] = int(line[j])
                index += 1
    return vector


def load_dataSet(data_dir):
    fileList = os.listdir(data_dir)
    x_data = []
    y_data = []
    for filename in fileList:
        x_data.append(vector(data_dir + filename))
        y_data.append(int(filename.split("_")[0]))
    return x_data, y_data


def KNN(test_data, train_data, k):
    labelList = []
    for i in test_data:
        distances = []
        for j in train_data:
            dis = np.sqrt(np.sum(np.square(i[:-1] - j[:-1])))
            distances.append([dis, j[-1]])
        distances.sort()
        distances_k = distances[:k]
        label_k = [y[1] for y in distances_k]
        label = max(label_k, key=label_k.count)
        labelList.append(label)
    test_y = test_data[:, -1]
    result = np.mean(test_y == labelList)
    return result


if __name__ == '__main__':
    train_dir = "trainingDigits/"
    test_dir = "testDigits/"

    trainX, trainY = load_dataSet(train_dir)
    testX, testY = load_dataSet(test_dir)

    train_data = np.c_[trainX, trainY]
    test_data = np.c_[testX, testY]

    acc = KNN(test_data, train_data, 3)
    print(acc)
